# 导入必要的库和模块
import gradio as gr
import numpy as np
import pickle
from PIL import Image

# 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
with open("best_knn_model.pkl", "rb") as file:
    best_knn_model = pickle.load(file)

# 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(image_data):
    img_array = np.array(image_data, dtype=np.uint8)
    # 将输入图像转换为灰度并调整大小
    img = Image.fromarray(img_array).convert('L').resize((8, 8))

    # 将图像转换为numpy数组并进行归一化
    img_array = np.array(img) // 16  # 缩放到0-15
    img_array = img_array.flatten().reshape(1, -1)  # 扁平化并添加批次维度

    # 使用加载的KNN模型进行预测
    prediction = best_knn_model.predict(img_array)
    return int(prediction[0])

# 创建Gradio接口，这个接口将用于用户输入和显示预测结果
interface = gr.Interface(
    fn=predict_digit,  # 传入预测函数
    inputs=gr.Sketchpad(),  # 输入为用户书写的图像
    outputs=gr.Label(),  # 输出为预测结果
    title="Handwritten Digit Recognition",
    description="Draw a digit from 0 to 9 on the sketchpad and let the model guess what it is."
)

# 启动Gradio接口，用户可以通过这个接口进行交互
interface.launch()